import numpy as np
import matplotlib.pyplot as plt
import torch
from mpl_toolkits.mplot3d import Axes3D
# z = (x^2 + y - 11)^2 + (x + y^2 -7）^2  找到局部极小值
def himmelblau(x):
    return (x[0]**2 + x[1] - 11)**2 + (x[0] + x[1]**2 - 7)**2

# 生成数据
x = np.arange(-6, 6, 0.1)
y = np.arange(-6, 6, 0.1)

X, Y = np.meshgrid(x, y)


Z = himmelblau([X, Y])

# 创建图形和三维轴对象
fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111, projection='3d')

# 绘制三维表面图
surface = ax.plot_surface(X, Y, Z, cmap='viridis', edgecolor='none')

# 设置视图角度
ax.view_init(elev=60, azim=-30)

# 设置轴标签
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')

# 添加颜色条
fig.colorbar(surface, shrink=0.5, aspect=5)

# 显示图形
plt.show()




# 初始化张量
x = torch.tensor([4., 0.], requires_grad=True)

# 定义优化器
optimizer = torch.optim.Adam([x], lr=1e-3)

for step in range(20000):
    optimizer.zero_grad()  # 清零梯度

    pred = himmelblau(x)  # 计算目标函数值

    pred.backward()  # 反向传播计算梯度

    optimizer.step()  # 更新参数

    if step % 2000 == 0:
        print(f'step {step}: x = {x.tolist()}, f(x) = {pred.item()}')
